import allure
import numpy as np
import torch
import torch.nn as nn


class RNN(nn.Module):
    def __init__(self, input_size, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.rnn = nn.GRU(
            input_size=input_size,
            hidden_size=64,
            num_layers=1,
            batch_first=True,
        )
        self.out = nn.Sequential(
            nn.Linear(64, 1),
        )
        self.hidden = None

    def forward(self, x):
        r_out, self.hidden = self.rnn(x)  # None 表示 hidden state 会用全0的 state
        out = self.out(r_out)
        return out


class LSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super(LSTM, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # x = x.view(len(x), 1, -1)  # 把原有2维度[a,b]改为3维[a,1,b]
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
        out = self.fc(out[:, -1, :])
        return out


def lstm_predict(x):
    # 输入的维度为1，只有Close收盘价
    input_dim = 1
    # 隐藏层特征的维度
    hidden_dim = 32
    # 循环的layers
    num_layers = 2
    # 预测后一天的收盘价
    output_dim = 1

    lstm = LSTM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers)
    lstm.load_state_dict(torch.load('model/lstm_model.bin'))

    x_hat = np.mean(x)
    x_std = np.std(x)
    input_data = (x - x_hat) / x_std
    output = lstm(torch.unsqueeze(torch.Tensor(input_data), dim=0).view(-1, len(input_data), 1))
    # output = lstm(torch.Tensor(input_data))

    output_value = torch.squeeze(output).cpu().detach().numpy()
    prediction = x_hat + output_value
    return prediction


def srl_predict(x):
    rnn = RNN(10)
    rnn.load_state_dict(torch.load('model/srl_model.bin'))
    x_hat = np.mean(x)
    x_std = np.std(x)
    input_data = (x - x_hat) / x_std
    output = rnn(torch.unsqueeze(torch.Tensor(input_data), dim=0))
    output_value = torch.squeeze(output).cpu().detach().numpy()
    prediction = x_hat + output_value
    return prediction


def nlp_predict(x):
    return x[-1]


@allure.feature("lstm模型测试")
def test_lstm():
    x = np.array([1650.3, 1646.3, 1633, 1654.2, 1620.3, 1616.3, 1611.1, 1622.2, 1646.3, 1633])
    res = lstm_predict(x)
    assert isinstance(res, float)


@allure.feature("srl模型测试")
def test_srl():
    x = np.array([1650.3, 1646.3, 1633, 1654.2, 1620.3, 1616.3, 1611.1, 1622.2, 1646.3, 1633])
    res = srl_predict(x)
    assert isinstance(res, float)


@allure.feature("nlp模型测试")
def test_nlp():
    x = np.array([1650.3, 1646.3, 1633, 1654.2, 1620.3, 1616.3, 1611.1, 1622.2, 1646.3, 1633])
    res = nlp_predict(x)
    assert isinstance(res, float)
